package net.hserver.watermark.util;

import cn.hutool.core.util.StrUtil;
import com.itextpdf.text.Element;
import com.itextpdf.text.Rectangle;
import com.itextpdf.text.pdf.*;
import org.apache.poi.ooxml.POIXMLDocumentPart;
import org.apache.poi.ooxml.POIXMLProperties;
import org.apache.poi.openxml4j.opc.PackagePartName;
import org.apache.poi.openxml4j.opc.PackageRelationship;
import org.apache.poi.openxml4j.opc.TargetMode;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xslf.usermodel.*;
import org.apache.poi.xssf.usermodel.XSSFRelation;
import org.apache.poi.xssf.usermodel.XSSFSheet;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.apache.poi.xwpf.model.XWPFHeaderFooterPolicy;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFHeader;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;
import org.apache.xmlbeans.XmlObject;
import org.openxmlformats.schemas.wordprocessingml.x2006.main.CTSectPr;

import javax.imageio.ImageIO;
import javax.xml.namespace.QName;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.*;
import java.util.HashMap;

public class WaterMarkUtils {

    public static boolean isOK(String fileName) {
        if (fileName.contains(".pdf")) {
            return true;
        } else if (fileName.contains(".jpg")) {
            return true;
        } else if (fileName.contains(".png")) {
            return true;
        } else if (fileName.contains(".docx") || fileName.contains(".doc")) {
            return true;
        } else if (fileName.contains(".xlsx") || fileName.contains(".xls")) {
            return true;
        } else if (fileName.contains(".ppt") || fileName.contains(".pptx")) {
            return true;
        }
        return false;
    }


    public static ByteArrayOutputStream addWaterMarkCommon(String waterMarkContent, String fileName, byte[] res) throws IOException {
        ByteArrayOutputStream sourceOs = new ByteArrayOutputStream();
        sourceOs.write(res);
        ByteArrayOutputStream targetOs;
        fileName = fileName.toLowerCase();
        ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(sourceOs.toByteArray());
        if (fileName.contains(".pdf")) {
            targetOs = WaterMarkUtils.markPdf(byteArrayInputStream, waterMarkContent);
        } else if (fileName.contains(".jpg")) {
            targetOs = WaterMarkUtils.markImage(byteArrayInputStream, waterMarkContent, "jpg");
        } else if (fileName.contains(".png")) {
            targetOs = WaterMarkUtils.markImage(byteArrayInputStream, waterMarkContent, "png");
        } else if (fileName.contains(".docx") || fileName.contains(".doc")) {
            targetOs = WaterMarkUtils.markWord(byteArrayInputStream, waterMarkContent);
        } else if (fileName.contains(".xlsx") || fileName.contains(".xls")) {
            targetOs = WaterMarkUtils.markExcel(byteArrayInputStream, waterMarkContent);
        } else if (fileName.contains(".ppt") || fileName.contains(".pptx")) {
            targetOs = WaterMarkUtils.markPPT(byteArrayInputStream, waterMarkContent);
        } else {
            targetOs = sourceOs;
        }
        return targetOs;
    }
//
//    public static void markVideo(String input, String output, String waterMarkContent, String format) {
//        FFmpegFrameGrabber frameGrabber;
//        Frame frame;
//        FFmpegFrameRecorder recorder;
//        try {
//            //抓取视频资源
//            frameGrabber = new FFmpegFrameGrabber(input);
//            frameGrabber.start();
//            recorder = toRecorder(format, frameGrabber, output);
//
//            recorder.start();
//            int count = 0;
//            int markCount = 0;
//            double randomValue = 0D;
//            System.gc();
//            while (true) {
//                frame = frameGrabber.grabFrame();
//                if (frame == null) {
//                    break;
//                }
//                //判断图片帧
//                if (frame.image != null) {
//                    if (count % 100 == 0) {
//                        markCount = 10;
//                        randomValue = Math.random();
//                    }
//                    if (markCount > 0) {
//                        markFrame(waterMarkContent, frame, recorder, randomValue);
//                        markCount--;
//                    } else {
//                        recorder.record(frame);
//                    }
//                }
//                //设置音频
//                if (frame.samples != null) {
//                    recorder.recordSamples(frame.sampleRate, frame.audioChannels, frame.samples);
//                }
//                count++;
//            }
//            closeAndDeleteData(frameGrabber, recorder);
//        } catch (Exception e) {
//            e.printStackTrace();
//        }
//    }
//
//    private static void closeAndDeleteData(FFmpegFrameGrabber frameGrabber, FFmpegFrameRecorder recorder) throws IOException {
//        recorder.stop();
//        recorder.release();
//        frameGrabber.stop();
//        System.gc();
//    }
//
//    private static FFmpegFrameRecorder toRecorder(String format, FFmpegFrameGrabber frameGrabber, String markFilePath) {
//        FFmpegFrameRecorder recorder;
//        recorder = new FFmpegFrameRecorder(markFilePath, frameGrabber.getImageWidth(), frameGrabber.getImageHeight(), frameGrabber.getAudioChannels());
//        recorder.setFormat(format);
//        recorder.setSampleRate(frameGrabber.getSampleRate());
//        recorder.setFrameRate(frameGrabber.getFrameRate());
//        recorder.setTimestamp(frameGrabber.getTimestamp());
//        recorder.setVideoBitrate(frameGrabber.getVideoBitrate());
//        recorder.setVideoCodec(frameGrabber.getVideoCodec());
//        return recorder;
//    }
//
//    private static void markFrame(String waterMarkContent, Frame frame, FFmpegFrameRecorder recorder, double randomValue) throws FrameRecorder.Exception {
//        IplImage iplImage = Java2DFrameUtils.toIplImage(frame);
//        BufferedImage buffImg = Java2DFrameUtils.toBufferedImage(iplImage);
//        Graphics2D graphics = buffImg.createGraphics();
//        graphics.setColor(Color.WHITE);
//        graphics.setComposite(AlphaComposite.getInstance(AlphaComposite.SRC_ATOP, 0.3f));
//        graphics.setFont(new Font("宋体", Font.BOLD, 20));
//        graphics.drawString(waterMarkContent, (int) (iplImage.width() * randomValue), (int) (iplImage.height() * randomValue));
//        graphics.dispose();
//        Frame newFrame = Java2DFrameUtils.toFrame(buffImg);
//        recorder.record(newFrame);
//        iplImage.close();
//    }

    /**
     * 给图片添加水印文字、可设置水印文字的大小、旋转角度、透明度
     *
     * @param logoText 水印内容
     */
    private static ByteArrayOutputStream markImage(ByteArrayInputStream inputStream, String logoText, String type) throws IOException {

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();


        // 1、读取源图片，Image获取图片宽度、高度
        Image scrImg = ImageIO.read(inputStream);
        BufferedImage buffImg = new BufferedImage(scrImg.getWidth(null), scrImg.getHeight(null), BufferedImage.TYPE_INT_RGB);

        // 2、得到画笔对象
        Graphics2D graphics = buffImg.createGraphics();

        // 3、设置对线段的锯齿状边缘处理
        graphics.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
        graphics.drawImage(scrImg.getScaledInstance(scrImg.getWidth(null), scrImg.getHeight(null), Image.SCALE_SMOOTH), 0, 0, null);

        // 4、设置水印倾斜度，这里是在图片的对角线上
        // 对角线长度lengthOfDiagonal
        double lengthOfDiagonal = Math.sqrt(Math.pow(buffImg.getWidth(), 2) + Math.pow(buffImg.getHeight(), 2));
        double v = (Math.pow(buffImg.getWidth(), 2) + Math.pow(lengthOfDiagonal, 2) - Math.pow(buffImg.getHeight(), 2)) / (2 * buffImg.getWidth() * lengthOfDiagonal);
        //get到了一个弧度数
        double acos = Math.acos(v);
        double myDegree = Math.toDegrees(acos);
        //这里的负号决定对角线-Math.toRadians(myDegree)
        graphics.rotate(-Math.toRadians(myDegree),
                (double) buffImg.getWidth() / 2,
                (double) buffImg.getHeight() / 2);

        // 5、设置水印文字颜色
        graphics.setColor(Color.DARK_GRAY);

        // 6、获取源图片的宽度和高度
        int width = scrImg.getWidth(null);
        int heigth = scrImg.getHeight(null);


        graphics.setFont(new Font("微软雅黑", Font.BOLD, 50));

        //8、设置透明度
        graphics.setComposite(AlphaComposite.getInstance(AlphaComposite.SRC_ATOP, 0.3f));

        //9、设置文字位置
//        FontDesignMetrics metrics = FontDesignMetrics.getMetrics(graphics.getFont());
        //获取文字宽度
        int strWidth = 100;

        int xNum = width / strWidth + 1;

        int yNum = heigth / 50 + 1;

        int split = 50;

        for (int i = 1; i <= 2 * yNum; i++) {
            int y = -heigth + 50 * i + 5 * split * i;
            for (int j = 0; j < xNum; j++) {
                int x = strWidth * j + 3 * split * j;

                graphics.drawString(logoText, x, y);
            }
        }

        //11、释放资源
        graphics.dispose();

        //12、生成图片
        try {
            ImageIO.write(buffImg, type, outputStream);
        } catch (
                FileNotFoundException e) {
            e.printStackTrace();
        }
        return outputStream;
    }

    /**
     * 给PDF添加水印
     *
     * @param inputStream      原文件流
     * @param waterMarkContent 添加水印的内容
     */
    private static ByteArrayOutputStream markPdf(ByteArrayInputStream inputStream, String waterMarkContent) {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try {
            // 水印的高和宽
            int waterMarkHeight = 30;
            int watermarkWeight = 60;

            // 水印间隔距离
            int waterMarkInterval = 100;

            // 读取PDF文件流
            PdfReader pdfReader = new PdfReader(inputStream);

            PdfDocument pdfDocument = new PdfDocument();

            // 创建PDF文件的模板，可以对模板的内容修改，重新生成新PDF文件
            PdfStamper pdfStamper = new PdfStamper(pdfReader, outputStream);

            //添加PDF自定义属性
            HashMap<String, String> info = pdfReader.getInfo();
            if (StrUtil.isEmpty(info.get("Title"))) {
                info.put("Title", waterMarkContent);
            }
            pdfStamper.setMoreInfo(info);

            // 设置水印字体
            BaseFont baseFont = BaseFont.createFont("STSong-Light", "UniGB-UCS2-H", BaseFont.EMBEDDED);

            // 设置PDF内容的Graphic State 图形状态
            PdfGState pdfGraPhicState = new PdfGState();
            // 填充透明度
            pdfGraPhicState.setFillOpacity(0.1f);
            // 轮廓不透明度
            pdfGraPhicState.setStrokeOpacity(0.1f);

            // PDF页数
            int pdfPageNum = pdfReader.getNumberOfPages() + 1;

            // PDF文件内容字节
            PdfContentByte pdfContent;

            // PDF页面矩形区域
            Rectangle pageRectangle;

            for (int i = 1; i < pdfPageNum; i++) {
                // 获取当前页面矩形区域
                pageRectangle = pdfReader.getPageSizeWithRotation(i);
                // 获取当前页内容，getOverContent表示之后会在页面内容的上方加水印
                pdfContent = pdfStamper.getOverContent(i);

                // 获取当前页内容，getOverContent表示之后会在页面内容的下方加水印
                // pdfContent = pdfStamper.getUnderContent(i);

                pdfContent.saveState();
                // 设置水印透明度
                pdfContent.setGState(pdfGraPhicState);

                // 开启写入文本
                pdfContent.beginText();
                // 设置字体
                pdfContent.setFontAndSize(baseFont, 20);

                // 在高度和宽度维度每隔waterMarkInterval距离添加一个水印
                for (int height = waterMarkHeight; height < pageRectangle.getHeight(); height = height + waterMarkInterval) {
                    for (int width = watermarkWeight; width < pageRectangle.getWidth() + watermarkWeight;
                         width = width + waterMarkInterval) {
                        // 添加水印文字并旋转30度角
                        pdfContent.showTextAligned(Element.ALIGN_LEFT, waterMarkContent, width - watermarkWeight,
                                height - waterMarkHeight, 30);
                    }
                }
                // 停止写入文本
                pdfContent.endText();
            }
            pdfStamper.close();
            pdfReader.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return outputStream;
    }

    private static ByteArrayOutputStream markWord(ByteArrayInputStream ins, String markStr) {
        ByteArrayOutputStream ous = new ByteArrayOutputStream();
        try {
            XWPFDocument doc = new XWPFDocument(ins);
            XWPFParagraph paragraph;
            CTSectPr sectPr = doc.getDocument().getBody().addNewSectPr();
            XWPFHeaderFooterPolicy headerFooterPolicy = new XWPFHeaderFooterPolicy(doc, sectPr);
            headerFooterPolicy.createWatermark(markStr);
            XWPFHeader header = headerFooterPolicy.getHeader(XWPFHeaderFooterPolicy.DEFAULT);
            paragraph = header.getParagraphArray(0);
            if (paragraph == null) {
                paragraph = doc.createParagraph();
            }

            XmlObject[] xmlobjects = paragraph.getCTP().getRArray(0).getPictArray(0).selectChildren(new QName("urn:schemas-microsoft-com:vml", "shape"));
            if (xmlobjects.length > 0) {
                com.microsoft.schemas.vml.CTShape ctshape = (com.microsoft.schemas.vml.CTShape) xmlobjects[0];
                // set fill color
                ctshape.setFillcolor("#d8d8d8");
                // set rotation
                ctshape.setStyle(ctshape.getStyle() + ";rotation:315");
            }
            doc.write(ous);
            doc.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return ous;
    }

    /**
     * 创建水印图片
     * excel
     *
     * @param waterMark 水印内容
     * @return
     */
    private static BufferedImage createWaterMarkImage(String waterMark) {
        String[] textArray = waterMark.split("\n");
        Font font = new Font("microsoft-yahei", Font.PLAIN, 20);
        Integer width = 500;
        Integer height = 200;
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
        // 背景透明 开始
        Graphics2D g = image.createGraphics();
        image = g.getDeviceConfiguration().createCompatibleImage(width, height, Transparency.TRANSLUCENT);
        g.dispose();
        // 背景透明 结束
        g = image.createGraphics();
        // 设定画笔颜色
        g.setColor(new Color(Integer.parseInt("#C5CBCF".substring(1), 16)));
        // 设置画笔字体
        g.setFont(font);
        // 设定倾斜度
        g.shear(0.1, -0.26);
        // 设置字体平滑
        g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        int y = 150;
        for (int i = 0; i < textArray.length; i++) {
            // 画出字符串
            g.drawString(textArray[i], 0, y);
            y = y + font.getSize();
        }
        // 释放画笔
        g.dispose();
        return image;
    }

    private static ByteArrayOutputStream markExcel(ByteArrayInputStream ins, String markStr) {
        ByteArrayOutputStream ous = new ByteArrayOutputStream();
        try {
            BufferedImage image = createWaterMarkImage(markStr);
            // 导出到字节流B
            ByteArrayOutputStream os = new ByteArrayOutputStream();
            ImageIO.write(image, "png", os);
            XSSFWorkbook workbook = new XSSFWorkbook(ins);
            //添加自定义属性
            POIXMLProperties.CustomProperties customProperties = workbook.getProperties().getCustomProperties();
            customProperties.addProperty("sample", markStr);

            int pictureIdx = workbook.addPicture(os.toByteArray(), Workbook.PICTURE_TYPE_PNG);
            POIXMLDocumentPart poixmlDocumentPart = workbook.getAllPictures().get(pictureIdx);

            //获取每个Sheet表
            for (int i = 0; i < workbook.getNumberOfSheets(); i++) {
                XSSFSheet sheet = workbook.getSheetAt(i);
                PackagePartName ppn = poixmlDocumentPart.getPackagePart().getPartName();
                String relType = XSSFRelation.IMAGES.getRelation();
                PackageRelationship pr = sheet.getPackagePart().addRelationship(ppn, TargetMode.INTERNAL, relType, null);
                sheet.getCTWorksheet().addNewPicture().setId(pr.getId());
            }
            workbook.write(ous);
            workbook.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return ous;
    }

    private static ByteArrayOutputStream markPPT(ByteArrayInputStream inputStream, String markStr) {
        ByteArrayOutputStream ous = new ByteArrayOutputStream();
        try {
            XMLSlideShow ppt = new XMLSlideShow(inputStream);
            //添加自定义属性
            POIXMLProperties.CustomProperties customProperties = ppt.getProperties().getCustomProperties();
            customProperties.addProperty("sample", markStr);
            Dimension pageSize = ppt.getPageSize();
            int fontSize = 80;
            int height = fontSize;
            int width = fontSize * (markStr.length() + 1);
            for (XSLFSlide slide : ppt.getSlides()) {
                // 在每一页上添加水印的代码 创建文本框对象
                XSLFTextShape waterMark = slide.createTextBox();
                waterMark.setTextRotation(-15D);
                XSLFTextParagraph paragraph = waterMark.addNewTextParagraph();
                XSLFTextRun run1 = paragraph.addNewTextRun();
                run1.setText(markStr);
                run1.setFontColor(Color.lightGray);
                run1.setFontSize(Double.parseDouble(fontSize + ""));
                run1.setFontFamily("宋体");
                paragraph.addLineBreak();
                // 设置水印位置和大小
                waterMark.setAnchor(new java.awt.Rectangle((int) pageSize.getWidth() / 2 - width / 2,
                        (int) pageSize.getHeight() / 2 - height / 2, width, height));
            }
            ppt.write(ous);
            ppt.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return ous;
    }
}

